import numpy as np
import random as rd

'''
思路:
    1、滤波器初始化，before = now = init_value
    2、将获得的[x,y,z]导入滤波器
    3、计算出速度[vx,vy,vz]
    4、将[x,y,z]和[vx,vy,vz]分为三个二维滤波器处理，即三个二维观测值
    



'''

time = 0.00222222
pre_time = 0.05


class KalmanFilter:
    '''

    对x，y，z三个方向 分别进行滤波

    '''

    def __init__(self):
        self.position_before_array = np.array([0.0000, 0.0000, 0.000])
        self.position_now_array = np.array([0.000, 0.0000, 0.0000])
        self.velocity = np.array([0.0000, 0.0000, 0.0000])
        self.flag = 1
        self.pos_v_array_now = np.zeros((2, 3))  # 三维测量值  z
        self.pos_v_array_predict = np.zeros((2, 3))  # 三维估计值  x
        self.K_x_y_z = []  # kalman gain 三个维度 K
        self.P_x_y_z = []  # 估计误差 三个维度 P
        for i in range(0, 3):  # init the k & p
            zero1 = np.random.random((2, 2))
            zero2 = np.random.random((2, 2))
            self.K_x_y_z.append(zero1)
            self.P_x_y_z.append(zero2)

        # parameter A , H , Q , R
        self.A = np.array([[1.0000, time], [0.0000, 1.0000]])
        self.H = np.array([[1.0000, 0.0000], [0.0000, 1.0000]])
        self.Q = np.array([[0.2500, 0.0000], [0.0000, 0.2500]])
        self.R = np.array([[0.0100, 0.0000], [0.0000, 0.01000]])
        self.I = np.array([[1.0000, 0.0000], [0, 1.0000]])
        self.pretime = np.array([[1.0000, pre_time], [0.0000, 1.0000]])
        self.pos_predict = np.array([0.0000, 0.0000, 0.0000])

    def process(self, array):  # array 是一个array类型
        '''
        向滤波器导入数据
        同时计算速度1264.
        生成三个二阶模型
        '''

        for i in range(0, 3):
            self.position_before_array[i] = self.position_now_array[i]
            self.position_now_array[i] = array[i]

        self.velocity = (self.position_now_array -
                         self.position_before_array)/time
        self.pos_v_array_now = np.vstack(
            (self.position_now_array, self.velocity))

        # print(self.pos_v_array_now)

        if self.flag == 1:
            self.pos_v_array_predict = self.pos_v_array_now
            self.flag = 0
        '''
        进行迭代处理三个维度
        '''
        for i in range(0, 3):
            x = self.pos_v_array_predict[:, i:i + 1]
            p = self.P_x_y_z[i]
            # 先验
            self.pos_v_array_predict[:, i:i + 1] = np.dot(self.A, x)  # 求x估计值
            self.P_x_y_z[i] = np.dot(
                np.dot(self.A, p), self.A.T)+self.Q  # 求估计误差协方差
            # 修正
            PHT = np.dot(p, self.H.T)
            HPHT_R = np.dot(np.dot(self.H, p), self.H.T)+self.R
            HPHT_R = np.linalg.inv(HPHT_R)
            self.K_x_y_z[i] = np.dot(PHT, HPHT_R)  # 修正K
            KH = np.dot(self.K_x_y_z[i], self.H)
            x_temp = np.dot(self.I-KH, self.pos_v_array_predict[:, i:i + 1])
            self.pos_v_array_predict[:, i:i + 1] = x_temp+np.dot(
                self.K_x_y_z[i], self.pos_v_array_now[:, i:i + 1])  # 修正x，得到后验
            self.P_x_y_z[i] = np.dot(self.I-KH, self.P_x_y_z[i])  # 修正p
        pos = np.zeros((2, 3))
        for i in range(0, 3):
            pos[:, i:i+1] = np.dot(self.pretime,
                                   self.pos_v_array_predict[:, i:i + 1])
        self.pos_predict = pos[0:1, :]
        output = self.pos_predict.tolist()
        return output[0]


if __name__ == '__main__':
    KF = KalmanFilter()
    pointarray = []
    for i in range(0, 12):
        pointarray.append([i*i/100, i*4/77, i/51])
    # print(pointarray)
    for i in range(0, 12):
        a = KF.process(pointarray[i])
        print("timed "+str(i)+" : ")
        print("real pos "+str(pointarray[i]))
        print("predict pos "+str(a))
